#include "image_process.h"
#include <iostream>
#include "coco_info.h"

using namespace std;


void preprocess(const cv::Mat& src, cv::Mat& dst)
{
    // Resize and pad image while meeting stride-multiple constraints
    double width = src.cols;
    double height = src.rows;
    // Scale ratio (new / old)
    double r = MIN(IMG_WIDTH / width, IMG_HEIGHT / height);
    
    // 图片中需要填充的像素数量
    int scale_width = (int)(std::round(width * r));
    int scale_height = (int)(std::round(height * r));
    int dw = (IMG_WIDTH - scale_width) / 2;
    int dh = (IMG_HEIGHT - scale_height) / 2;

    cv::resize(src, dst, cv::Size(scale_width, scale_height), 0.0, 0.0, cv::INTER_LINEAR);
    int top = (int)(std::round(dh - 0.1));
    int bottom = (int)(std::round(dh + 0.1));
    int left = (int)(std::round(dw - 0.1));
    int right = (int)(std::round(dw + 0.1));

    cv::copyMakeBorder(dst, dst, top, bottom, left, right, cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114));
}

void generate_bboxes(const int8_t* bsnn_out, const float* anchors, int out_h, int out_w, int stride,
                     float dequant_scale, std::vector<ObjInfo> &objects)
{
    float tmp = 0.0;
    for (int y = 0; y < out_h; y++)
    {
        for (int x = 0; x < out_w; x++)
        {
            auto start_pos = ((y * out_w) + x) * OUT_CHANNELS;
            for (int q = 0; q < 3; q++)
            {
                float anchor_w = anchors[q * 2];
                float anchor_h = anchors[q * 2 + 1];
                auto idx = start_pos + q * PER_ANCHORS_CHN;

                //find class index with max score
                int class_index = 0;
                int8_t class_score_ = std::numeric_limits<int8_t>::lowest();
                for (int k = 0; k < NUM_CLASSES; k++)
                {
                    int8_t score = bsnn_out[idx + 5 + k];
                    if (score > class_score_)
                    {
                        class_score_ = score;
                        class_index = k + 1;
                    }
                }

                float class_score = class_score_ * dequant_scale;
                float box_score = bsnn_out[idx + 4] * dequant_scale;
                float confidence = sigmoid(box_score) * sigmoid(class_score);
                confidence = sigmoid(confidence);
                
                if (confidence > tmp)
                {
                    tmp = confidence;
                }
                if (confidence >= CONF_THRESHOLD)
                {
                    float dx = sigmoid(bsnn_out[idx]      * dequant_scale);
                    float dy = sigmoid(bsnn_out[idx + 1]  * dequant_scale);
                    float dw = sigmoid(bsnn_out[idx + 2]  * dequant_scale);
                    float dh = sigmoid(bsnn_out[idx + 3]  * dequant_scale);

                    float pb_cx = (dx * 2.f -0.5f + x) * stride;
                    float pb_cy = (dy * 2.f - 0.5f + y) * stride;
                    float pb_w = pow(dw * 2.f, 2) * anchor_w;
                    float pb_h = pow(dh * 2.f, 2) * anchor_h;

                    float x0 = pb_cx - pb_w * 0.5f;
                    float y0 = pb_cy - pb_h * 0.5f;
                    float x1 = pb_cx + pb_w * 0.5f;
                    float y1 = pb_cy + pb_h * 0.5f;
                    ObjInfo obj;
                    obj.classId = class_index;
                    obj.score = confidence;
                    obj.x1 = x0;
                    obj.y1 = y0;
                    obj.x2 = x1;
                    obj.y2 = y1;
                    objects.emplace_back(obj);
                }
            }
        }
    }

    cout << tmp << endl;

    
}

void process_output(bsnn_user_output_t *bsnn_output, std::vector<ObjInfo> &result)
{
    int strides[] = {8, 16, 32};
    std::vector<ObjInfo> det_objs;
    // stride 8
    {
        float anchors[] = {10.f, 13.f, 16.f, 30.f, 33.f, 23.f};
        int out_h = IMG_HEIGHT / strides[0];
        int out_w = IMG_WIDTH / strides[0];
        int8_t ts_out_i8[out_h * out_w * OUT_CHANNELS] =  {0};

        // DWH格式下，在（i， j）位置包含D个数据，将i, j拉平和进行处理
        int8_t* out_buf_i8 = (int8_t*)(bsnn_output->ts_addr[0]);
        for (int i = 0; i < out_h * out_w; i++)
        {
            for (int j = 0; j < OUT_CHANNELS; j++)
            {
                ts_out_i8[i * OUT_CHANNELS + j] = out_buf_i8[i * REG_PITCH + j];
            }
        }

        generate_bboxes(ts_out_i8, anchors, out_h, out_w, strides[0],
                        bsnn_output->ts_qnt_scale[0], det_objs);
    }

    // stride 16
    {
        float anchors[] = {30.f, 61.f, 62.f, 45.f, 59.f, 119.f};
        int out_h = IMG_HEIGHT / strides[1];
        int out_w = IMG_WIDTH / strides[1];
        int8_t ts_out_i8[out_h * out_w * OUT_CHANNELS] =  {0};

        // DWH格式下，在（i， j）位置包含D个数据，将i, j拉平和进行处理
        int8_t* out_buf_i8 = (int8_t*)(bsnn_output->ts_addr[1]);
        for (int i = 0; i < out_h * out_w; i++)
        {
            for (int j = 0; j < OUT_CHANNELS; j++)
            {
                ts_out_i8[i * OUT_CHANNELS + j] = out_buf_i8[i * REG_PITCH + j];
            }
        }

        generate_bboxes(ts_out_i8, anchors, out_h, out_w, strides[1],
                        bsnn_output->ts_qnt_scale[1], det_objs);
    }

    // stride 32
    {
        float anchors[] = {116.f, 90.f, 156.f, 198.f, 373.f, 326.f};
        int out_h = IMG_HEIGHT / strides[2];
        int out_w = IMG_WIDTH / strides[2];
        int8_t ts_out_i8[out_h * out_w * OUT_CHANNELS] =  {0};

        // DWH格式下，在（i， j）位置包含D个数据，将i, j拉平和进行处理
        int8_t* out_buf_i8 = (int8_t*)(bsnn_output->ts_addr[2]);
        for (int i = 0; i < out_h * out_w; i++)
        {
            for (int j = 0; j < OUT_CHANNELS; j++)
            {
                ts_out_i8[i * OUT_CHANNELS + j] = out_buf_i8[i * REG_PITCH + j];
            }
        }

        generate_bboxes(ts_out_i8, anchors, out_h, out_w, strides[2],
                        bsnn_output->ts_qnt_scale[2], det_objs);
    }

    // cout << det_objs.size() << endl;
    cout << "-> The number of objects detected: " << det_objs.size() << endl;

    soft_nms(det_objs, result, 2, 0.5, CONF_THRESHOLD , NMS_THRESHOLD);


}



void draw_bboxes(cv::Mat& img, const std::vector<ObjInfo>& result)
{
    float width = img.cols;
    float height = img.rows;
    float gain = MIN(IMG_WIDTH / width, IMG_HEIGHT / height);
    float pad_w = (IMG_WIDTH - width * gain) / 2;
    float pad_h = (IMG_HEIGHT - height * gain) / 2;
    
    stringstream str;
    for (int i = 0; i <  result.size(); i++)
    {
        int cls_id = result[i].classId;
        std::string cls_name = coco_classes[cls_id];
        cv::Scalar color = coco_colors[cls_id];
        str << result[i].score;
        std::string confidence = str.str().substr(0, 3);
        int x1 = (int)((result[i].x1 - pad_w) / gain);
        int y1 = (int)((result[i].y1 - pad_h) / gain);
        int x2 = (int)((result[i].x2 - pad_w) / gain);
        int y2 = (int)((result[i].y2 - pad_h) / gain);

        cv::rectangle(img, cv::Rect(cv::Point2i(x1, y1), cv::Point2i(x2, y2)), color, 2);
        string label = cls_name + " " + confidence;
        cv::putText(img, label, cv::Point2i(x1 + 10, y1 - 10), cv::FONT_HERSHEY_SIMPLEX, 0.7, color, 1, cv::LINE_AA);


    }
}